import logging
import torch
import numpy as np
from numpy.linalg import norm
import itertools
from crowd_sim.envs.policy.policy import Policy
from crowd_sim.envs.utils.action import ActionRot, ActionXY
from crowd_sim.envs.utils.state import tensor_to_joint_state
from crowd_sim.envs.utils.utils import point_to_segment_dist
from crowd_nav.policy.state_predictor import StatePredictor, LinearStatePredictor_batch
from crowd_nav.policy.graph_model import Social_Stgcnn
from crowd_nav.policy.value_estimator import DQNNetwork
from crowd_nav.policy.reward_estimate import estimate_reward_on_predictor


class MpStgcnn(Policy):
    def __init__(self):
        super().__init__()
        self.name = 'MpStgcnn'
        self.trainable = True
        self.multiagent_training = True
        self.kinematics = None
        self.epsilon = None
        self.gamma = None
        self.sampling = None
        self.speed_samples = None
        self.rotation_samples = None
        self.action_space = None
        self.rotation_constraint = None
        self.speeds = None
        self.rotations = None
        self.action_values = None
        self.robot_state_dim = 9
        self.human_state_dim = 5
        self.v_pref = 1
        self.share_graph_model = None
        self.value_estimator = None
        self.linear_state_predictor = None
        self.state_predictor = None
        self.planning_depth = None
        self.planning_width = None
        self.do_action_clip = None
        self.sparse_search = None
        self.sparse_speed_samples = 2
        self.sparse_rotation_samples = 8
        self.action_group_index = []
        self.traj = None
        self.use_noisy_net = False
        self.count=0

    def configure(self, config, device):
        self.set_common_parameters(config)
        self.planning_depth = config.mp_stgcnn.planning_depth
        self.do_action_clip = config.mp_stgcnn.do_action_clip
        if hasattr(config.mp_stgcnn, 'sparse_search'):
            self.sparse_search = config.mp_stgcnn.sparse_search
        self.planning_width = config.mp_stgcnn.planning_width
        self.share_graph_model = config.mp_stgcnn.share_graph_model
        self.linear_state_predictor = config.mp_stgcnn.linear_state_predictor
        # self.set_device(device)
        self.device = device

        graph_model1 = Social_Stgcnn(config, self.robot_state_dim, self.human_state_dim)
        self.value_estimator = DQNNetwork(config, graph_model1)
        graph_model2 = Social_Stgcnn(config, self.robot_state_dim, self.human_state_dim)
        self.state_predictor = StatePredictor(config, graph_model2, self.time_step)
        self.model = [graph_model1, graph_model2, self.value_estimator.value_network,
                      self.state_predictor.human_motion_predictor]

        logging.info('Planning depth: {}'.format(self.planning_depth))
        logging.info('Planning width: {}'.format(self.planning_width))
        logging.info('Sparse search: {}'.format(self.sparse_search))

        if self.planning_depth > 1 and not self.do_action_clip:
            logging.warning('Performing d-step planning without action space clipping!')

    def set_common_parameters(self, config):
        self.gamma = config.rl.gamma
        self.kinematics = config.action_space.kinematics
        self.sampling = config.action_space.sampling
        self.speed_samples = config.action_space.speed_samples
        self.rotation_samples = config.action_space.rotation_samples
        self.rotation_constraint = config.action_space.rotation_constraint

    def set_device(self, device):
        self.device = device
        for model in self.model:
            model.to(device)

    def set_epsilon(self, epsilon):
        self.epsilon = epsilon

    def set_noisy_net(self, use_noisy_net):
        self.use_noisy_net = use_noisy_net

    def set_time_step(self, time_step):
        self.time_step = time_step
        self.state_predictor.time_step = time_step

    def get_normalized_gamma(self):
        return pow(self.gamma, self.time_step * self.v_pref)

    def get_model(self):
        return self.value_estimator

    def get_state_dict(self):
        if self.state_predictor.trainable:
            if self.share_graph_model:
                return {
                    'graph_model': self.value_estimator.graph_model.state_dict(),
                    'value_network': self.value_estimator.value_network.state_dict(),
                    'motion_predictor': self.state_predictor.human_motion_predictor.state_dict()
                }
            else:
                return {
                    'graph_model1': self.value_estimator.graph_model.state_dict(),
                    'graph_model2': self.state_predictor.graph_model.state_dict(),
                    'value_network': self.value_estimator.value_network.state_dict(),
                    'motion_predictor': self.state_predictor.human_motion_predictor.state_dict()
                }
        else:
            return {
                    'graph_model': self.value_estimator.graph_model.state_dict(),
                    'value_network': self.value_estimator.value_network.state_dict()
                }

    def get_traj(self):
        return self.traj

    def load_state_dict(self, state_dict):
        if self.state_predictor.trainable:
            if self.share_graph_model:
                self.value_estimator.graph_model.load_state_dict(state_dict['graph_model'])
            else:
                self.value_estimator.graph_model.load_state_dict(state_dict['graph_model1'])
                self.state_predictor.graph_model.load_state_dict(state_dict['graph_model2'])

            self.value_estimator.value_network.load_state_dict(state_dict['value_network'])
            self.state_predictor.human_motion_predictor.load_state_dict(state_dict['motion_predictor'])
        else:
            self.value_estimator.graph_model.load_state_dict(state_dict['graph_model'])
            self.value_estimator.value_network.load_state_dict(state_dict['value_network'])

    def save_model(self, file):
        torch.save(self.get_state_dict(), file)

    def load_model(self, file):
        checkpoint = torch.load(file)
        self.load_state_dict(checkpoint)

    def build_action_space(self, v_pref):
        """
        Action space consists of 25 uniformly sampled actions in permitted range and 25 randomly sampled actions.
        """
        holonomic = True if self.kinematics == 'holonomic' else False
        # speeds = [(np.exp((i + 1) / self.speed_samples) - 1) / (np.e - 1) * v_pref for i in range(self.speed_samples)]
        speeds = [(i+1)/self.speed_samples * v_pref for i in range(self.speed_samples)]
        if holonomic:
            rotations = np.linspace(0, 2 * np.pi, self.rotation_samples, endpoint=False)
        else:
            rotations = np.linspace(-self.rotation_constraint, self.rotation_constraint, self.rotation_samples)

        action_space = [ActionXY(0, 0) if holonomic else ActionRot(0, 0)]
        self.action_group_index.append(0)

        for i, rotation in enumerate(rotations):
            for j, speed in enumerate(speeds):
                action_index = i * self.speed_samples + j + 1
                self.action_group_index.append(action_index)
                if holonomic:
                    action_space.append(ActionXY(speed * np.cos(rotation), speed * np.sin(rotation)))
                else:
                    action_space.append(ActionRot(speed, rotation))
        self.speeds = speeds
        self.rotations = rotations
        self.action_space = action_space

    def predict(self, state):
        """
        A base class for all methods that takes pairwise joint state as input to value network.
        The input to the value network is always of shape (batch_size, # humans, rotated joint state length)

        """
        self.count=self.count+1
        if self.phase is None or self.device is None:
            raise AttributeError('Phase, device attributes have to be set!')
        if self.phase == 'train' and self.epsilon is None:
            raise AttributeError('Epsilon attribute has to be set in training phase')

        if self.reach_destination(state):
            return ActionXY(0, 0) if self.kinematics == 'holonomic' else ActionRot(0, 0)
        if self.action_space is None:
            self.build_action_space(1.0)
        max_action = None
        origin_max_value = float('-inf')
        state_tensor = state.to_tensor(add_batch_size=True, device=self.device)


        probability = np.random.random()
        if self.phase == 'train' and probability < self.epsilon and self.use_noisy_net is False:
            max_action_index = np.random.choice(len(self.action_space))
            max_action = self.action_space[max_action_index]
            self.last_state = self.transform(state)
            return max_action, max_action_index
        else:
            max_value, max_action_index, max_traj = self.V_planning(state_tensor)
            if max_value[0] > origin_max_value:
                max_action = self.action_space[max_action_index[0]]
            if max_action is None:
                raise ValueError('Value network is not well trained.')

        if self.phase == 'train':
            self.last_state = self.transform(state)
        else:
            self.last_state = self.transform(state)
            self.traj = max_traj[0]
        return max_action, int(max_action_index[0])

    def V_planning(self, state):
        """ Plans n steps into future based on state action value function. Computes the value for the current state as well as the trajectories
        defined as a list of (state, action, reward) triples
        """
        # current_state_value = self.value_estimator(state)
        robot_state_batch = state[0]
        human_state_batch = state[1]
        q_value = torch.Tensor(self.value_estimator(state)[0].to('cpu'))
        max_action_value, max_action_indexes = torch.max(q_value, dim=1)
        trajs = []
        for i in range(robot_state_batch.shape[0]):
            cur_state = (robot_state_batch[i, :, :].unsqueeze(0), human_state_batch[i, :, :].unsqueeze(0))
            trajs.append([(cur_state, None, None)])
        return max_action_value, max_action_indexes, trajs



    # def transform(self, state):
    #     """
    #     Take the JointState to tensors
    #
    #     :param state:
    #     :return: tensor of shape (# of agent, len(state))
    #     """
    #     robot_state_tensor = torch.Tensor([state.robot_state.to_tuple()]).to(self.device)
    #     human_states_tensor = torch.Tensor([human_state.to_tuple() for human_state in state.human_states]). \
    #         to(self.device)
    #
    #     return robot_state_tensor, human_states_tensor
    def transform(self, state_stack):
        """
        Take the JointState to tensors

        :param state:
        :return: tensor of shape (# of agent, len(state))
        """
        robot_state_tensor = torch.Tensor([[state.to_tuple() for state in state_stack.robot_state_stack]]).to(self.device)
        robot_state_tensor = robot_state_tensor.permute(1,0,2)
        human_states_tensor = torch.Tensor([[human_state.to_tuple() for human_state in human_states]
                                            for human_states in state_stack.human_states_stack]).to(self.device)

        return robot_state_tensor, human_states_tensor

    def compute_next_robot_state(self, robot_state, action):
        if robot_state.shape[0] != 1:
            raise NotImplementedError
        next_state = robot_state.clone().squeeze()
        if self.kinematics == 'holonomic':
            next_state[0] = next_state[0] + action.vx * self.time_step
            next_state[1] = next_state[1] + action.vy * self.time_step
            next_state[2] = action.vx
            next_state[3] = action.vy
        else:
            next_state[7] = next_state[7] + action.r
            next_state[0] = next_state[0] + np.cos(next_state[7]) * action.v * self.time_step
            next_state[1] = next_state[1] + np.sin(next_state[7]) * action.v * self.time_step
            next_state[2] = np.cos(next_state[7]) * action.v
            next_state[3] = np.sin(next_state[7]) * action.v
        return next_state.unsqueeze(0).unsqueeze(0)
    def get_attention_weights(self):
        return self.value_estimator.graph_model.attention_weights

    @staticmethod
    def reach_destination(state_stack):
        robot_state_stack = state_stack.robot_state_stack
        if np.linalg.norm((robot_state_stack[-1].py - robot_state_stack[-1].gy,
                           robot_state_stack[-1].px - robot_state_stack[-1].gx)) < robot_state_stack[-1].radius:
            return True
        else:
            return False